'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 1_study_fewshot_prompt_template.py
* @Time: 2025/9/10
* @All Rights Reserve By Brtc
'''
import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()
# 1、构建示例模板与示例
example = [
    {"question":"帮我计算一下2+2等于多少？", "answer":"4"},
    {"question":"帮我计算一下2+3等于多少？", "answer":"5"},
    {"question":"帮我计算一下2*15等于多少？", "answer":"300"}
]
exapmle_prompt = ChatPromptTemplate.from_messages([
    ("human", "{question}"),
    ("ai", "{answer}")
])
#2、构建少量提示词模板
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=exapmle_prompt,
    examples=example
)
#3、构建最终提示词模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个可以计算复杂数学问题的聊天机器人，请严格按照示例样本输出， 不要废话！"),
    few_shot_prompt,
    ("human","{question}")
])
#4、创建大语言模型与链
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
chain = prompt | llm | StrOutputParser()
#5、调用链获取结果
print(chain.invoke("帮我计算5*7等于多少？"))